import os
# 设置环境变量，允许重复加载OpenMP库（临时 workaround）
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
from ultralytics import YOLO

model = YOLO("./yolov8n.pt")  # 加载预训练模型（建议用于训练）

if __name__ == '__main__':
    #imgsz=640输入大小
    model.train(data="./dataset/mydata.yaml", epochs=3, imgsz=640)  # 训练模型

    metrics = model.val()  # 在验证集上评估模型性能
    success = model.export(format="onnx")  # 将模型导出为 ONNX 格式
    results = model("./dataset/test/images/00002.png")  # 对图像进行预测
    # Process results list
    for result in results:
        boxes = result.boxes  # Boxes object for bounding box outputs
        masks = result.masks  # Masks object for segmentation masks outputs
        keypoints = result.keypoints  # Keypoints object for pose outputs
        probs = result.probs  # Probs object for classification outputs
        result.show()  # display to screen